import heapq
import pandas as pd
import numpy as np
import random
import matplotlib.pyplot as plt
from collections import defaultdict
import warnings

warnings.filterwarnings('ignore')
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False

# 从Excel文件读取罗马尼亚地图和启发函数数据 - 复用work1的代码
def read_data_from_excel(file_path):
    print("开始从Excel文件读取数据...")
    
    try:
        # 读取第一个工作表(sheet1)中的距离矩阵
        print("正在从Sheet1读取邻接矩阵...")
        distance_df = pd.read_excel(file_path, sheet_name=0, header=0, index_col=0)
        print(f"邻接矩阵读取成功! 形状: {distance_df.shape}")
        
        # 读取第二个工作表(sheet2)中的启发函数值
        print("正在从Sheet2读取启发函数值...")
        heuristic_df = pd.read_excel(file_path, sheet_name=1)
        print(f"启发函数值读取成功! 形状: {heuristic_df.shape}")
        
        return distance_df, heuristic_df
    
    except Exception as e:
        print(f"读取Excel数据失败: {str(e)}")
        raise

# 从数据框构建罗马尼亚地图 - 复用work1的代码
def create_romania_map(distance_df):
    graph = defaultdict(list)
    print("开始构建罗马尼亚地图...")
    connections_count = 0
    
    # 根据距离矩阵构建邻接表
    for city1 in distance_df.index:
        for city2 in distance_df.columns:
            # 获取距离值并确保它是一个标量
            try:
                distance = distance_df.loc[city1, city2]
                # 使用标量比较方法而不是直接使用 != 运算符
                if isinstance(distance, (int, float)):
                    if distance != 1000:
                        graph[city1].append((city2, distance))
                        connections_count += 1
                else:
                    # 如果是Series，转换为标量后比较
                    distance_val = distance.iloc[0] if hasattr(distance, 'iloc') else float(distance)
                    if distance_val != 1000:
                        graph[city1].append((city2, distance_val))
                        connections_count += 1
            except Exception as e:
                print(f"处理城市 {city1} 到 {city2} 的连接时出错: {e}")
    
    print(f"罗马尼亚地图构建完成! 共有 {len(graph)} 个城市, {connections_count} 条连接")
    return graph

# 构建到布加勒斯特的启发函数 - 复用work1的代码
def build_heuristic_function(heuristic_df):
    heuristic = {}
    print("开始构建启发函数...")
    
    # 确定列名 (假定第一列是城市名，第二列是启发函数值)
    if len(heuristic_df.columns) < 2:
        print(f"警告: 启发函数数据格式不符合预期，列数为 {len(heuristic_df.columns)}")
    
    # 读取城市名和启发函数值
    count = 0
    for i in range(len(heuristic_df)):
        try:
            city = heuristic_df.iloc[i, 0]  # 第一列是城市名
            value = heuristic_df.iloc[i, 1]  # 第二列是启发函数值
            if pd.notna(city) and pd.notna(value):
                heuristic[city] = value
                count += 1
        except Exception as e:
            print(f"处理启发函数的第 {i} 行时出错: {e}")
    
    print(f"启发函数构建完成! 共加载了 {count} 个城市的启发值")
    
    # 打印部分启发函数值作为调试信息
    if heuristic:
        print("部分启发函数值示例:")
        items = list(heuristic.items())[:5]
        for city, value in items:
            print(f"  {city}: {value}")
    
    return heuristic

# 创建城市索引映射，方便蚁群算法处理
def create_city_indices(romania_map):
    cities = list(romania_map.keys())
    city_to_idx = {city: i for i, city in enumerate(cities)}
    idx_to_city = {i: city for i, city in enumerate(cities)}
    return cities, city_to_idx, idx_to_city

# 创建距离矩阵 - 用于蚁群算法
def create_distance_matrix(romania_map, city_to_idx):
    n = len(city_to_idx)
    # 初始化为很大的值
    distance_matrix = np.ones((n, n)) * 1000
    
    # 填充实际距离
    for city, neighbors in romania_map.items():
        i = city_to_idx[city]
        for neighbor, distance in neighbors:
            j = city_to_idx[neighbor]
            distance_matrix[i, j] = distance
    
    return distance_matrix

# 蚁群算法实现
class AntColonyOptimization:
    def __init__(self, distance_matrix, n_ants=10, alpha=1.0, beta=2.0, 
                 rho=0.5, Q=100, max_iterations=100):
        """
        初始化蚁群算法参数
        
        参数:
        - distance_matrix: 城市之间的距离矩阵
        - n_ants: 蚂蚁数量
        - alpha: 信息素重要性参数
        - beta: 启发式信息重要性参数
        - rho: 信息素蒸发率
        - Q: 信息素增加常数
        - max_iterations: 最大迭代次数
        """
        self.distance_matrix = distance_matrix
        self.n_cities = distance_matrix.shape[0]
        self.n_ants = n_ants
        self.alpha = alpha
        self.beta = beta
        self.rho = rho
        self.Q = Q
        self.max_iterations = max_iterations
        
        # 初始化信息素矩阵
        self.pheromone = np.ones((self.n_cities, self.n_cities))
        # 计算启发式信息矩阵（距离的倒数）
        self.heuristic = 1.0 / (distance_matrix + np.eye(self.n_cities))
        
        # 存储最优路径
        self.best_path = None
        self.best_path_length = float('inf')
        self.convergence_history = []
    
    def run(self, start_idx, end_idx):
        """
        运行蚁群算法
        
        参数:
        - start_idx: 起始城市索引
        - end_idx: 目标城市索引
        
        返回:
        - best_path: 最优路径的城市索引列表
        - best_path_length: 最优路径长度
        """
        print(f"启动蚁群算法，从城市索引 {start_idx} 到 {end_idx}")
        
        for iteration in range(self.max_iterations):
            # 构建所有蚂蚁的路径
            all_paths = []
            all_path_lengths = []
            
            for ant in range(self.n_ants):
                path, path_length = self._construct_path(start_idx, end_idx)
                all_paths.append(path)
                all_path_lengths.append(path_length)
                
                # 更新最优路径
                if path_length < self.best_path_length:
                    self.best_path = path
                    self.best_path_length = path_length
                    print(f"迭代 {iteration+1}: 发现新的最优路径，长度为 {path_length}")
            
            # 信息素蒸发
            self.pheromone *= (1 - self.rho)
            
            # 信息素更新
            for path, path_length in zip(all_paths, all_path_lengths):
                self._update_pheromone(path, path_length)
                
            # 记录收敛历史
            self.convergence_history.append(self.best_path_length)
            
            # 打印当前迭代的信息
            if (iteration + 1) % 10 == 0:
                print(f"迭代 {iteration+1}/{self.max_iterations}, 当前最优路径长度: {self.best_path_length}")
        
        return self.best_path, self.best_path_length
    
    def _construct_path(self, start_idx, end_idx):
        """为一只蚂蚁构建从起点到终点的路径"""
        current = start_idx
        path = [current]
        path_length = 0.0
        visited = {current}
        
        # 继续构建路径直到到达目标城市
        while current != end_idx:
            next_city = self._select_next_city(current, visited)
            
            # 如果没有可访问的城市，终止路径构建
            if next_city is None:
                break
                
            # 更新路径和路径长度
            path.append(next_city)
            path_length += self.distance_matrix[current, next_city]
            visited.add(next_city)
            current = next_city
            
            # 如果到达目标城市，结束路径构建
            if current == end_idx:
                break
        
        return path, path_length
    
    def _select_next_city(self, current, visited):
        """
        根据概率选择下一个城市
        
        参数:
        - current: 当前城市索引
        - visited: 已访问城市集合
        
        返回:
        - 下一个城市的索引，如果没有可访问城市则返回None
        """
        # 计算未访问城市的选择概率
        unvisited = [i for i in range(self.n_cities) if i not in visited 
                    and self.distance_matrix[current, i] < 1000]
        
        if not unvisited:
            return None
        
        # 计算概率
        pheromone = np.array([self.pheromone[current, j] for j in unvisited])
        heuristic = np.array([self.heuristic[current, j] for j in unvisited])
        probabilities = (pheromone ** self.alpha) * (heuristic ** self.beta)
        
        # 归一化概率
        probabilities = probabilities / probabilities.sum()
        
        # 根据概率选择下一个城市
        next_idx = np.random.choice(len(unvisited), p=probabilities)
        return unvisited[next_idx]
    
    def _update_pheromone(self, path, path_length):
        """
        更新路径上的信息素
        
        参数:
        - path: 路径上的城市索引列表
        - path_length: 路径长度
        """
        # 计算信息素增量
        delta = self.Q / path_length
        
        # 更新路径上的信息素
        for i in range(len(path) - 1):
            city1, city2 = path[i], path[i + 1]
            self.pheromone[city1, city2] += delta
            self.pheromone[city2, city1] += delta  # 对称更新
    
    def plot_convergence(self):
        """绘制收敛曲线"""
        plt.figure(figsize=(10, 6))
        plt.plot(range(1, len(self.convergence_history) + 1), self.convergence_history, marker='o')
        plt.title('蚁群算法收敛曲线')
        plt.xlabel('迭代次数')
        plt.ylabel('最优路径长度')
        plt.grid(True)
        plt.savefig('aco_convergence.png')
        plt.show()

def ant_colony_search(romania_map, start_city, goal_city, city_to_idx, idx_to_city, 
                     n_ants=20, alpha=1.0, beta=5.0, rho=0.5, max_iterations=100):
    """
    使用蚁群算法解决罗马尼亚度假问题
    
    参数:
    - romania_map: 罗马尼亚地图(邻接表)
    - start_city: 起始城市名称
    - goal_city: 目标城市名称
    - city_to_idx: 城市名称到索引的映射
    - idx_to_city: 索引到城市名称的映射
    
    返回:
    - path: 最优路径(城市名称列表)
    - path_length: 最优路径长度
    - aco: 蚁群算法实例(用于获取更多信息)
    """
    # 创建距离矩阵
    distance_matrix = create_distance_matrix(romania_map, city_to_idx)
    
    # 获取起点和终点的索引
    start_idx = city_to_idx[start_city]
    goal_idx = city_to_idx[goal_city]
    
    # 初始化蚁群算法
    aco = AntColonyOptimization(
        distance_matrix, 
        n_ants=n_ants, 
        alpha=alpha, 
        beta=beta, 
        rho=rho,
        max_iterations=max_iterations
    )
    
    # 运行算法
    best_path_indices, best_path_length = aco.run(start_idx, goal_idx)
    
    # 将路径索引转换为城市名称
    best_path = [idx_to_city[idx] for idx in best_path_indices]
    
    return best_path, best_path_length, aco

def main():
    # Excel文件路径
    file_path = r"罗马尼亚度假问题地图及启发函数值.xlsx"
    
    try:
        # 读取Excel数据
        distance_df, heuristic_df = read_data_from_excel(file_path)
        
        # 添加调试信息
        print("\n距离矩阵的形状:", distance_df.shape)
        print("距离矩阵的前几行:")
        print(distance_df.head())
        
        print("\n启发函数数据的形状:", heuristic_df.shape)
        print("启发函数数据的前几行:")
        print(heuristic_df.head())
        
        # 构建罗马尼亚地图
        romania_map = create_romania_map(distance_df)
        
        # 构建启发函数
        heuristic = build_heuristic_function(heuristic_df)
        
        # 创建城市索引映射
        cities, city_to_idx, idx_to_city = create_city_indices(romania_map)
        
        start_city = "Arad"
        goal_city = "Bucharest"
        
        # 验证起点和终点在图中
        print(f"\n验证起点和终点...")
        if start_city not in romania_map:
            print(f"错误: 起点 '{start_city}' 不在地图中!")
            return
        else:
            print(f"起点 '{start_city}' 验证成功!")
            
        if goal_city not in romania_map:
            print(f"错误: 终点 '{goal_city}' 不在地图中!")
            return
        else:
            print(f"终点 '{goal_city}' 验证成功!")
        
        print("\n罗马尼亚度假问题 - 从{}到{}的路径搜索".format(start_city, goal_city))
        print("=" * 50)
        
        # 打印地图信息
        print("\n地图信息:")
        print("城市数量:", len(romania_map))
        print("连接数量:", sum(len(neighbors) for neighbors in romania_map.values()))
        
        # 使用蚁群算法
        print("\n执行蚁群算法...")
        
        # 蚁群算法参数
        n_ants = 20           # 蚂蚁数量
        alpha = 1.0           # 信息素重要性
        beta = 5.0            # 启发式信息重要性
        rho = 0.5             # 信息素蒸发率
        max_iterations = 100  # 最大迭代次数
        
        best_path, best_path_length, aco = ant_colony_search(
            romania_map, start_city, goal_city, city_to_idx, idx_to_city,
            n_ants=n_ants, alpha=alpha, beta=beta, rho=rho, max_iterations=max_iterations
        )
        
        print("\n蚁群算法搜索结果:")
        if best_path:
            print("路径:", " -> ".join(best_path))
            print("总代价:", best_path_length)
            print("算法参数:")
            print(f"  蚂蚁数量: {n_ants}")
            print(f"  信息素重要性(alpha): {alpha}")
            print(f"  启发式信息重要性(beta): {beta}")
            print(f"  信息素蒸发率(rho): {rho}")
            print(f"  最大迭代次数: {max_iterations}")
            
            # 绘制收敛曲线
            aco.plot_convergence()
        else:
            print("未找到路径!")
            
    except Exception as e:
        print("执行过程中发生错误:", str(e))
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()